import networkx as nx
from collections import defaultdict, deque
from llvmlite import ir


class Mem2Reg:
    def __init__(self, function):
        self.function = function
        self.dom_tree = None
        self.df = None
        self.alloca_info = {}
        self.phi_placements = defaultdict(set)
        self.value_stacks = defaultdict(list)
        self.phi_nodes = {}
        self.alloca_users = defaultdict(list)  # 存储 alloca 的用户指令

    def run(self):
        # 步骤1: 构建CFG和支配树
        self.build_cfg()
        self.compute_dominance()

        # 步骤2: 收集所有指令的使用关系
        self.collect_instruction_uses()

        # 步骤3: 识别可提升的alloca指令
        self.collect_promotable_allocas()

        # 步骤4: 计算支配边界并插入Phi节点
        self.compute_phi_placements()
        self.insert_phi_nodes()

        # 步骤5: 变量重命名
        self.rename_variables()

        # 步骤6: 清理旧指令
        self.cleanup()

    def collect_instruction_uses(self):
        """收集所有指令的使用关系"""
        for block in self.function.blocks:
            for inst in block.instructions:
                for operand in inst.operands:
                    if isinstance(operand, ir.Instruction):
                        # 记录指令被哪些指令使用
                        self.alloca_users[operand].append(inst)

    def build_cfg(self):
        """构建控制流图(CFG)"""
        self.cfg = nx.DiGraph()
        self.block_map = {}

        # 添加基本块到CFG
        for block in self.function.blocks:
            self.block_map[block.name] = block
            self.cfg.add_node(block)

        # 添加边(基于终结指令)
        for block in self.function.blocks:
            terminator = block.terminator
            if terminator is None:
                continue

            if terminator.opname == "br":
                if len(terminator.operands) == 1:
                    # 无条件跳转
                    succ = terminator.operands[0]
                    self.cfg.add_edge(block, succ)
                else:
                    # 条件跳转
                    _, true_succ, false_succ = terminator.operands
                    self.cfg.add_edge(block, true_succ)
                    self.cfg.add_edge(block, false_succ)
            elif terminator.opname == "ret":
                # 返回指令，无后继
                pass
            # 可以添加其他终结指令处理

    def compute_dominance(self):
        """计算支配树和支配边界"""
        if not self.cfg.nodes:
            return

        # 使用NetworkX计算支配树
        entry = next(iter(self.function.blocks))
        self.dom_tree = nx.immediate_dominators(self.cfg, entry)

        # 计算支配边界
        self.df = defaultdict(set)
        for node in self.cfg.nodes:
            if self.cfg.in_degree(node) > 1:
                for predecessor in self.cfg.predecessors(node):
                    runner = predecessor
                    while runner != self.dom_tree.get(node, None) and runner != entry:
                        self.df[runner].add(node)
                        runner = self.dom_tree.get(runner, None)

    def collect_promotable_allocas(self):
        """收集可提升的alloca指令"""
        entry_block = next(iter(self.function.blocks))
        for inst in entry_block.instructions:
            if isinstance(inst, ir.AllocaInstr):
                # 检查是否只被load/store使用
                promotable = True
                users = self.alloca_users.get(inst, [])

                for user in users:
                    if not isinstance(user, (ir.LoadInstr, ir.StoreInstr)):
                        promotable = False
                        break

                if promotable:
                    self.alloca_info[inst] = {"def_blocks": set(), "users": users}

                    # 收集定义块
                    for user in users:
                        if isinstance(user, ir.StoreInstr) and user.operands[1] == inst:
                            self.alloca_info[inst]["def_blocks"].add(user.parent)

    def compute_phi_placements(self):
        """计算需要插入Phi节点的位置"""
        for alloca, info in self.alloca_info.items():
            # 使用迭代支配边界算法
            worklist = list(info["def_blocks"])
            visited = set()

            while worklist:
                block = worklist.pop()

                if block not in self.df:
                    continue

                for frontier in self.df[block]:
                    if frontier not in visited:
                        visited.add(frontier)
                        self.phi_placements[alloca].add(frontier)
                        worklist.append(frontier)

    def insert_phi_nodes(self):
        """插入Phi节点"""
        for alloca, blocks in self.phi_placements.items():
            for block in blocks:
                # 在块开头创建Phi节点
                phi = ir.PhiInstr(block, alloca.type.pointee, name=f"{alloca.name}.phi")

                # 在第一条指令前插入Phi节点
                if block.instructions:
                    block.instructions.insert(0, phi)
                else:
                    block.instructions.append(phi)

                self.phi_nodes[(alloca, block)] = phi

    def rename_variables(self):
        """变量重命名"""
        # 初始化值栈
        for alloca in self.alloca_info:
            self.value_stacks[alloca] = [ir.Undefined]

        # 深度优先遍历支配树
        entry = next(iter(self.function.blocks))
        self.rename_block(entry)

    def rename_block(self, block):
        """重命名块中的变量"""
        # 保存当前栈状态（用于恢复）
        stack_tops = {a: len(stack) for a, stack in self.value_stacks.items()}

        # 处理Phi节点
        for inst in block.instructions:
            if isinstance(inst, ir.PhiInstr):
                for (alloca, b), phi in self.phi_nodes.items():
                    if b == block and phi == inst:
                        # 将Phi节点压入栈
                        self.value_stacks[alloca].append(inst)

        # 处理其他指令
        for inst in list(block.instructions):  # 使用副本遍历
            # 处理Load指令
            if isinstance(inst, ir.LoadInstr):
                ptr = inst.operands[0]
                if ptr in self.value_stacks and self.value_stacks[ptr]:
                    # 用栈顶值替换Load
                    new_value = self.value_stacks[ptr][-1]

                    # 替换所有使用
                    for user in list(self.alloca_users.get(inst, [])):
                        for i, op in enumerate(user.operands):
                            if op == inst:
                                user.operands = (
                                    list(user.operands[:i])
                                    + [new_value]
                                    + list(user.operands[i + 1 :])
                                )

                    # 从块中删除Load指令
                    if inst in block.instructions:
                        block.instructions.remove(inst)

            # 处理Store指令
            elif isinstance(inst, ir.StoreInstr):
                value, ptr = inst.operands
                if ptr in self.value_stacks:
                    # 更新值栈
                    self.value_stacks[ptr].append(value)

                    # 从块中删除Store指令
                    if inst in block.instructions:
                        block.instructions.remove(inst)

        # 为后继块的Phi节点添加输入
        for succ in self.cfg.successors(block):
            for inst in succ.instructions:
                if isinstance(inst, ir.PhiInstr):
                    for (alloca, b), phi in self.phi_nodes.items():
                        if b == succ and phi == inst and alloca in self.value_stacks:
                            # 添加当前值作为输入
                            phi.add_incoming(self.value_stacks[alloca][-1], block)

        # 递归处理支配树中的子节点
        for child in self.cfg.successors(block):
            if self.dom_tree.get(child) == block:
                self.rename_block(child)

        # 恢复栈状态
        for alloca, top in stack_tops.items():
            self.value_stacks[alloca] = self.value_stacks[alloca][:top]

    def cleanup(self):
        """清理旧指令"""
        # 删除alloca指令
        for alloca in self.alloca_info:
            if alloca in alloca.parent.instructions:
                alloca.parent.instructions.remove(alloca)

        # 删除剩余的store指令
        for alloca, info in self.alloca_info.items():
            for user in list(info["users"]):
                if isinstance(user, ir.StoreInstr) and user in user.parent.instructions:
                    user.parent.instructions.remove(user)


def create_complex_function(apply_mem2reg=False):
    # 创建模块和函数
    module = ir.Module()
    func_type = ir.FunctionType(ir.IntType(32), [ir.IntType(32)])
    func = ir.Function(module, func_type, name="complex_func")

    # 创建基本块
    entry = func.append_basic_block(name="entry")
    loop_cond = func.append_basic_block(name="loop_cond")
    loop_body = func.append_basic_block(name="loop_body")
    after_loop = func.append_basic_block(name="after_loop")
    fib_cond = func.append_basic_block(name="fib_cond")
    fib_body = func.append_basic_block(name="fib_body")
    fib_after = func.append_basic_block(name="fib_after")
    ret_block = func.append_basic_block(name="return")

    # 构建entry块
    builder = ir.IRBuilder(entry)
    n = func.args[0]

    # 分配变量
    n_addr = builder.alloca(ir.IntType(32), name="n.addr")
    fact_addr = builder.alloca(ir.IntType(32), name="fact.addr")
    i_addr = builder.alloca(ir.IntType(32), name="i.addr")
    a_addr = builder.alloca(ir.IntType(32), name="a.addr")
    b_addr = builder.alloca(ir.IntType(32), name="b.addr")
    temp_addr = builder.alloca(ir.IntType(32), name="temp.addr")
    fib_addr = builder.alloca(ir.IntType(32), name="fib.addr")
    j_addr = builder.alloca(ir.IntType(32), name="j.addr")
    result_addr = builder.alloca(ir.IntType(32), name="result.addr")

    # 初始化变量
    builder.store(n, n_addr)
    builder.store(ir.Constant(ir.IntType(32), 1), fact_addr)
    builder.store(ir.Constant(ir.IntType(32), 1), i_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), a_addr)
    builder.store(ir.Constant(ir.IntType(32), 1), b_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), fib_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), j_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), result_addr)

    # 跳转到循环条件
    builder.branch(loop_cond)

    # 构建循环条件块
    builder = ir.IRBuilder(loop_cond)
    i_val = builder.load(i_addr, name="i.val")
    n_val = builder.load(n_addr, name="n.val")
    cmp = builder.icmp_signed("<=", i_val, n_val, name="cmp")
    builder.cbranch(cmp, loop_body, after_loop)

    # 构建循环体块
    builder = ir.IRBuilder(loop_body)
    fact_val = builder.load(fact_addr, name="fact.val")
    i_val = builder.load(i_addr, name="i.val")
    new_fact = builder.mul(fact_val, i_val, name="new.fact")
    builder.store(new_fact, fact_addr)

    new_i = builder.add(i_val, ir.Constant(ir.IntType(32), 1), name="new.i")
    builder.store(new_i, i_addr)

    builder.branch(loop_cond)

    # 构建循环后块
    builder = ir.IRBuilder(after_loop)
    fact_val = builder.load(fact_addr, name="fact.final")
    builder.store(fact_val, result_addr)

    # 初始化斐波那契变量
    builder.store(ir.Constant(ir.IntType(32), 0), j_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), a_addr)
    builder.store(ir.Constant(ir.IntType(32), 1), b_addr)

    builder.branch(fib_cond)

    # 构建斐波那契条件块
    builder = ir.IRBuilder(fib_cond)
    j_val = builder.load(j_addr, name="j.val")
    n_val = builder.load(n_addr, name="n.val")
    cmp_fib = builder.icmp_signed("<", j_val, n_val, name="cmp.fib")
    builder.cbranch(cmp_fib, fib_body, fib_after)

    # 构建斐波那契体块
    builder = ir.IRBuilder(fib_body)
    a_val = builder.load(a_addr, name="a.val")
    b_val = builder.load(b_addr, name="b.val")

    # 计算下一个斐波那契数
    next_fib = builder.add(a_val, b_val, name="next.fib")
    builder.store(next_fib, fib_addr)

    # 更新变量
    builder.store(b_val, a_addr)
    builder.store(next_fib, b_addr)

    # 增加计数器
    new_j = builder.add(j_val, ir.Constant(ir.IntType(32), 1), name="new.j")
    builder.store(new_j, j_addr)

    builder.branch(fib_cond)

    # 构建斐波那契后块
    builder = ir.IRBuilder(fib_after)
    fib_val = builder.load(fib_addr, name="fib.val")
    result_val = builder.load(result_addr, name="result.val")
    final_result = builder.add(result_val, fib_val, name="final.result")
    builder.store(final_result, result_addr)

    # 跳转到返回块
    builder.branch(ret_block)

    # 构建返回块
    builder = ir.IRBuilder(ret_block)
    result_val = builder.load(result_addr, name="result.final")
    builder.ret(result_val)

    # 应用 mem2reg 优化
    if apply_mem2reg:
        mem2reg = Mem2Reg(func)
        mem2reg.run()
    
    return module


def create_square_sum_function(apply_mem2reg=False):
    # 创建模块和函数
    module = ir.Module()
    func_type = ir.FunctionType(ir.IntType(32), [ir.IntType(32)])
    func = ir.Function(module, func_type, name="square_sum")
    
    # 创建基本块
    entry = func.append_basic_block(name="entry")
    loop_cond = func.append_basic_block(name="loop_cond")
    loop_body = func.append_basic_block(name="loop_body")
    exit_block = func.append_basic_block(name="exit")
    
    # 构建entry块
    builder = ir.IRBuilder(entry)
    n = func.args[0]
    
    # 分配变量
    n_addr = builder.alloca(ir.IntType(32), name="n.addr")
    sum_addr = builder.alloca(ir.IntType(32), name="sum.addr")
    i_addr = builder.alloca(ir.IntType(32), name="i.addr")
    
    # 初始化变量
    builder.store(n, n_addr)
    builder.store(ir.Constant(ir.IntType(32), 0), sum_addr)  # sum = 0
    builder.store(ir.Constant(ir.IntType(32), 1), i_addr)     # i = 1
    
    builder.branch(loop_cond)
    
    # 构建循环条件块
    builder = ir.IRBuilder(loop_cond)
    i_val = builder.load(i_addr, name="i.val")
    n_val = builder.load(n_addr, name="n.val")
    cmp = builder.icmp_signed('<=', i_val, n_val, name="cmp")
    builder.cbranch(cmp, loop_body, exit_block)
    
    # 构建循环体块
    builder = ir.IRBuilder(loop_body)
    i_val = builder.load(i_addr, name="i.val")
    square = builder.mul(i_val, i_val, name="square")
    
    sum_val = builder.load(sum_addr, name="sum.val")
    new_sum = builder.add(sum_val, square, name="new.sum")
    builder.store(new_sum, sum_addr)
    
    new_i = builder.add(i_val, ir.Constant(ir.IntType(32), 1), name="new.i")
    builder.store(new_i, i_addr)
    
    builder.branch(loop_cond)
    
    # 构建退出块
    builder = ir.IRBuilder(exit_block)
    sum_val = builder.load(sum_addr, name="sum.final")
    builder.ret(sum_val)
    
    # 应用 mem2reg 优化
    if apply_mem2reg:
        mem2reg = Mem2Reg(func)
        mem2reg.run()
    
    return module

def compile_and_run(module, n):
    """编译并运行LLVM模块中的square_sum函数"""
    from ctypes import CFUNCTYPE, c_int
    import llvmlite.binding as llvm
    
    # 初始化LLVM环境
    llvm.initialize()
    llvm.initialize_native_target()
    llvm.initialize_native_asmprinter()
    
    def create_execution_engine():
        """创建JIT执行引擎"""
        target = llvm.Target.from_default_triple()
        target_machine = target.create_target_machine()
        backing_mod = llvm.parse_assembly("")
        engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
        return engine
    
    # 1. 获取LLVM IR字符串
    llvm_ir = str(module)
    
    # 2. 创建执行引擎
    engine = create_execution_engine()
    
    # 3. 编译IR
    mod = llvm.parse_assembly(llvm_ir)
    mod.verify()  # 验证模块
    
    # 4. 添加到引擎
    engine.add_module(mod)
    engine.finalize_object()
    
    # 5. 获取函数地址
    func_ptr = engine.get_function_address("complex_func")
    
    # 6. 通过ctypes调用
    func_type = CFUNCTYPE(c_int, c_int)
    func = func_type(func_ptr)
    result = func(n)
    
    return result

# 测试函数
def test_mem2reg():
    # 创建未优化的模块
    unoptimized_module = create_complex_function(apply_mem2reg=False)
    
    # 创建优化后的模块
    optimized_module = create_complex_function(apply_mem2reg=True)
    
    # 打印未优化的IR
    print("=== Unoptimized IR ===")
    print(unoptimized_module)
    
    # 打印优化后的IR
    print("\n=== Optimized IR (after mem2reg) ===")
    print(optimized_module)
    
    # 测试不同输入值
    test_values = [0, 1, 5, 10, 100]
    
    print("\n=== Test Results ===")
    print(f"{'n':>5} | {'Expected':>10} | {'Unoptimized':>12} | {'Optimized':>12} | {'Match':>5}")
    print("-" * 60)
    
    for n in test_values:
        # 计算预期结果
        expected = sum(i*i for i in range(1, n+1))
        
        # 运行未优化版本
        unoptimized_result = compile_and_run(unoptimized_module, n)

        # 运行优化版本
        optimized_result = compile_and_run(optimized_module, n)
        
        # 检查结果是否一致
        match = "✅" if unoptimized_result == optimized_result == expected else "❌"
        
        print(f"{n:>5} | {expected:>10} | {unoptimized_result:>12} | {optimized_result:>12} | {match:>5}")

# 运行测试
if __name__ == "__main__":
    test_mem2reg()